package com.ikas.ai.server.common;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import lombok.Data;

import java.util.Objects;

@Data
public class UpdateWorkConditionFlowJsonUtil {

    private static final String UPDATE_POLICY = "update_policy";
    private static final String EARLY_CACHE = "early_cache";
    private static final String FROM_STEP = "from_step";
    private static final String STEP0 = "step0";
    private static final String STEP1 = "step1";

    public static String updateWorkConditionFlowJson(String config, Integer updatePolicy, Integer step01EarlyCache, Integer stepNot01EarlyCache) throws JsonProcessingException {
        ObjectMapper mapper = new ObjectMapper();
        JsonNode root = mapper.readTree(config);
        ArrayNode arrayNode = (ArrayNode) root;
        JsonNode fromStep;
        for (JsonNode element : arrayNode) {
            fromStep = element.get(FROM_STEP);
            if (Objects.nonNull(fromStep)) {
                if (STEP1.equals(fromStep.asText()) || STEP0.equals(fromStep.asText())) {
                    modifyEarlyCacheStep1(element, step01EarlyCache);
                }
            } else {
                modifyEarlyCache(element, stepNot01EarlyCache);
            }
        }
        // 修改所有的 update_policy 为 2
        modifyUpdatePolicy(root, updatePolicy);
        // 将修改后的 JSON 树转换回 JSON 字符串
        return mapper.writeValueAsString(root);
    }

    private static void modifyUpdatePolicy(JsonNode node, int newValue) {
        if (node.isArray()) {
            ArrayNode arrayNode = (ArrayNode) node;
            for (JsonNode element : arrayNode) {
                modifyUpdatePolicy(element, newValue);
            }
        } else if (node.isObject()) {
            ObjectNode objectNode = (ObjectNode) node;
            if (objectNode.has(UPDATE_POLICY)) {
                objectNode.put(UPDATE_POLICY, newValue);
            }
            objectNode.fields().forEachRemaining(entry -> modifyUpdatePolicy(entry.getValue(), newValue));
        }
    }

    private static void modifyEarlyCacheStep1(JsonNode node, int newValue) {
        if (node.isArray()) {
            ArrayNode arrayNode = (ArrayNode) node;
            for (JsonNode element : arrayNode) {
                modifyEarlyCacheStep1(element, newValue);
            }
        }
        if (node.isObject()) {
            ObjectNode objectNode = (ObjectNode) node;
            if (objectNode.has(EARLY_CACHE)) {
                objectNode.put(EARLY_CACHE, newValue);
            }
            objectNode.fields().forEachRemaining(entry -> modifyEarlyCacheStep1(entry.getValue(), newValue));
        }
    }

    private static void modifyEarlyCache(JsonNode node, int newValue) {
        if (node.isArray()) {
            ArrayNode arrayNode = (ArrayNode) node;
            for (JsonNode element : arrayNode) {
                modifyEarlyCache(element, newValue);
            }
        } else if (node.isObject()) {
            ObjectNode objectNode = (ObjectNode) node;
            if (objectNode.has(EARLY_CACHE)) {
                objectNode.put(EARLY_CACHE, newValue);
            }
            objectNode.fields().forEachRemaining(entry -> modifyEarlyCache(entry.getValue(), newValue));
        }
    }
}